import matplotlib.pyplot as plt
import numpy as np
font = {'family': 'Arial', 'size' : 16}
plt.rc('font', **font)

plt.rcParams['mathtext.fontset'] = 'custom'
plt.rcParams['mathtext.it'] = 'Arial:italic'
plt.rcParams['mathtext.rm'] = 'Arial'
plt.rcParams['pdf.fonttype'] = 42

# Fixing random state for reproducibility
np.random.seed(19680801)

x = np.array([0, 1, 2, 3])
x_reverse = np.array([6, 5, 4, 3])
# linear noise
data0 = np.zeros((8, 4))
data0[0] = [0.40307, 0.40391, 0.40816, 0.40085] # no noise onion
data0[1] = [0.55491, 0.48641, 0.42181, 0.40654] # low noise onion
data0[2] = [0.85933, 0.64513, 0.55322, 0.47442] # medium noise onion
data0[3] = [1.14253, 0.84505, 0.67054, 0.58012] # high noise onion
data0[4] = [0.45477, 0.46012, 0.4584, 0.44659] # no noise 1by1
data0[5] = [0.63781, 0.55083, 0.49189, 0.46893] # low noise 1by1
data0[6] = [0.94354, 0.75398, 0.56962, 0.47317] # medium noise 1by1
data0[7] = [1.30899, 0.96826, 0.66788, 0.52819] # high noise 1by1

# normal noise
data1 = np.zeros((8, 4))
data1[0] = [0.40307, 0.40391, 0.40816, 0.40085] # no noise onion
data1[1] = [0.60644, 0.51568, 0.46783, 0.44573] # low noise onion
data1[2] = [0.80749, 0.68255, 0.59879, 0.55761] # medium noise onion
data1[3] = [0.92668, 0.82353, 0.70587, 0.6232] # high noise onion
data1[4] = [0.45477, 0.46012, 0.4584, 0.44659] # no noise 1by1
data1[5] = [0.69394, 0.61118, 0.53117, 0.5202] # low noise 1by1 
data1[6] = [0.9649, 0.82114, 0.68907, 0.59197] # medium noise 1by1
data1[7] = [1.04579, 1.00155, 0.85024, 0.71255] # high noise 1by1

fig, axs = plt.subplots(1, 2, constrained_layout=True, sharey=True, figsize=(11, 5))


axs[0].set_ylabel("MAE (eV)")

axs[0].text(0.3, 1.25, '\'1by1\' on \nlinear noise')
axs[0].text(3.2, 1.25, '\'onion\' on \nlinear noise')
axs[1].text(0.2, 1.25, '\'1by1\' on \nsampled noise')
axs[1].text(3.2, 1.25, '\'onion\' on \nsampled noise')

axs[0].annotate("(a)", xy=(0, 0.5), xytext=(-0.4, 1.3))
axs[1].annotate("(b)", xy=(0, 0.5), xytext=(-0.4, 1.3))
for ax in axs.flat:
    ax.set_xlabel("Step #")
    ax.set_xlim(0, 6)
    ax.set_xticks([0, 1, 2, 3, 4, 5, 6], labels=['0', '1', '2', '3', '2', '1', '0'])
    # ax.set_yticks([0.5, 1.5, 2.5, 3.5], labels=['None', 'Low', 'Med', 'High'])

for i in range(4):
    axs[0].quiver(x[:-1], data0[4+i][:-1], x[1:]-x[:-1], data0[4+i][1:]-data0[4+i][:-1], scale_units='xy', angles='xy', scale=1, 
    color='blue', width=0.0001, linewidth=0.5*i+1, edgecolor='black', headwidth=200, headlength=400)
    axs[0].quiver(x_reverse[:-1], data0[i][:-1], x_reverse[1:]-x_reverse[:-1], data0[i][1:]-data0[i][:-1], scale_units='xy', angles='xy', scale=1, 
    color='red', linestyle=':', linewidth=0.5*i+1, edgecolor='black', facecolor='none', width=0.0001, headwidth=200, headlength=400)
    axs[1].quiver(x[:-1], data1[4+i][:-1], x[1:]-x[:-1], data1[4+i][1:]-data1[4+i][:-1], scale_units='xy', angles='xy', scale=1, 
    color='blue', width=0.0001, linewidth=0.5*i+1, edgecolor='black', headwidth=200, headlength=400)
    axs[1].quiver(x_reverse[:-1], data1[i][:-1], x_reverse[1:]-x_reverse[:-1], data1[i][1:]-data1[i][:-1], scale_units='xy', angles='xy', scale=1, 
    color='red', linestyle=':', linewidth=0.5*i+1, edgecolor='black', facecolor='none', width=0.0001, headwidth=200, headlength=400)

for ax in axs.flat:
    ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[1])
    ax.plot([3, 3], [ax.get_ylim()[0], ax.get_ylim()[1]], '--', color="Grey")
    # ax.grid()


plt.show()